from flask import Flask, request, jsonify
from flask_cors import CORS
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import json
from datetime import datetime
import pymysql
import uuid
import os
from dotenv import load_dotenv
from concurrent.futures import ThreadPoolExecutor
import threading
import time

# Load environment variables
load_dotenv()

app = Flask(__name__)
CORS(app)  # Enable CORS for all routes

# 创建线程池
executor = ThreadPoolExecutor(max_workers=10)

# Configuration
API_URL = "http://8.136.111.93:8008/v1/chat-messages"
API_KEY = "app-iVySNIkK8TF0MoKc7GPbCJPG"
USER_ID = "test-user-001"
WORKFLOW_INPUT_KEY = "query"

# 创建带有重试机制的session
session = requests.Session()
retry_strategy = Retry(
    total=3,  # 总重试次数
    backoff_factor=1,  # 重试间隔
    status_forcelist=[500, 502, 503, 504]  # 需要重试的HTTP状态码
)
adapter = HTTPAdapter(max_retries=retry_strategy)
session.mount("http://", adapter)
session.mount("https://", adapter)

# Database configuration
DB_CONFIG = {
    'host': os.getenv('DB_HOST', 'localhost'),
    'user': os.getenv('DB_USER', 'root'),
    'password': os.getenv('DB_PASSWORD', '123456'),
    'db': os.getenv('DB_NAME', 'agricultureaigclargemodel'),
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}

def get_db_connection():
    try:
        return pymysql.connect(**DB_CONFIG)
    except pymysql.Error as e:
        print(f"Database connection error: {str(e)}")
        return None

def format_time():
    now = datetime.now()
    return f"{now.year}年{now.month}月{now.day}日 {now.hour}:{str(now.minute).zfill(2)}"

def save_chat_history(conversation_id, user_id, question, answer):
    conn = None
    try:
        conn = get_db_connection()
        if not conn:
            print("Failed to connect to database")
            return
            
        with conn.cursor() as cursor:
            chat_id = f"chat_{int(datetime.now().timestamp())}_{uuid.uuid4().hex[:8]}"
            chat_time = datetime.now()
            
            sql = """INSERT INTO chat_history 
                    (conversation_id, chat_id, user_id, chat_time, user_input, bot_response) 
                    VALUES (%s, %s, %s, %s, %s, %s)"""
            cursor.execute(sql, (conversation_id, chat_id, user_id, chat_time, question, answer))
            conn.commit()
    except Exception as e:
        print(f"Error saving chat history: {str(e)}")
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()

def call_ai_api(question, conversation_id, user_id, max_retries=3):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "query": question,
        "inputs": {WORKFLOW_INPUT_KEY: question},
        "response_mode": "blocking",
        "user": user_id
    }
    
    if conversation_id:
        payload["conversation_id"] = conversation_id

    for attempt in range(max_retries):
        try:
            # 使用session进行请求，它会自动处理重试
            response = session.post(
                API_URL,
                headers=headers,
                json=payload,
                timeout=(5, 60)  # (连接超时, 读取超时)
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"API call attempt {attempt + 1} failed: {str(e)}")
            if attempt == max_retries - 1:  # 最后一次尝试
                return None
            time.sleep(2 ** attempt)  # 指数退避
    return None

@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.json
    question = data.get('question')
    conversation_id = data.get('conversation_id')
    user_id = data.get('user_id', USER_ID)
    
    if not question:
        return jsonify({'error': 'No question provided'}), 400

    try:
        # 调用AI API获取响应
        result = call_ai_api(question, conversation_id, user_id)
        if not result:
            return jsonify({
                'error': 'AI服务暂时不可用，请稍后重试',
                'type': 'error',
                'content': '抱歉，服务器响应超时，请稍后重试。',
                'time': format_time()
            }), 500

        answer = result.get('answer', '')
        conv_id = result.get('conversation_id')
        
        # 异步保存聊天历史
        executor.submit(save_chat_history, conv_id, user_id, question, answer)
        
        return jsonify({
            'type': 'response',
            'content': answer,
            'time': format_time(),
            'conversation_id': conv_id,
            'relatedQuestions': [
                '如何选择具有抗病虫害特性的水稻品种？',
                '如何评估水稻品种的抗病虫害性？',
                '通过分析方法以及如何防范改良品种的抗病虫害性？'
            ]
        })

    except Exception as e:
        print(f"Chat error: {str(e)}")
        return jsonify({
            'error': str(e),
            'type': 'error',
            'content': '抱歉，发生了一些错误，请稍后重试。',
            'time': format_time()
        }), 500

@app.route('/api/chat/history', methods=['GET'])
def get_chat_history():
    user_id = request.args.get('user_id', USER_ID)
    conversation_id = request.args.get('conversation_id')
    page = int(request.args.get('page', 1))
    page_size = int(request.args.get('pageSize', 10))
    
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            where_clause = ["1=1"]
            params = []
            
            if user_id:
                where_clause.append("user_id = %s")
                params.append(user_id)
            
            if conversation_id:
                where_clause.append("conversation_id = %s")
                params.append(conversation_id)
            
            count_sql = f"SELECT COUNT(*) as total FROM chat_history WHERE {' AND '.join(where_clause)}"
            cursor.execute(count_sql, params)
            total = cursor.fetchone()['total']
            
            offset = (page - 1) * page_size
            sql = f"""
                SELECT * FROM chat_history 
                WHERE {' AND '.join(where_clause)}
                ORDER BY chat_time DESC
                LIMIT %s OFFSET %s
            """
            cursor.execute(sql, params + [page_size, offset])
            results = cursor.fetchall()
            
            return jsonify({
                'total': total,
                'rows': results,
                'page': page,
                'pageSize': page_size
            })
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        if 'conn' in locals():
            conn.close()

if __name__ == '__main__':
    app.run(debug=True, port=5000) 